﻿# coding=utf-8
# 2023/4/27: 首个记录，基于旧2020/11/3版本
# 2023/5/5: 枚举基于Enum
# 2023/8/17: 新增瞬时雨量字段。道路名称按新方式输出
# 2023/10/10: 补充部分英文注释
# 2023/10/24: 新增字段（累计雨量、精确风向、总辐射强度、紫外线辐射强度、臭氧浓度、一氧化碳浓度、二氧化硫浓度、二氧化氮浓度）
# 2025/3/26: 新增环境数据插值功能
# 2025/7/18: 修正significant赋值
# 2025/7/21: 使用类型注解。优化to_general_sample。修正speed limit插值

import bi_common as bi
from enum import Enum
from typing import List, Any


class WeatherType(Enum):
    UNKNOWN = 0  # 未知
    SUNNY = 1  # 晴
    CLOUDY = 2  # 阴
    RAINY = 3  # 雨
    SNOWY = 4  # 雪
    FOGGY = 5  # 雾
    SAND = 6  # 沙尘


class RoadType(Enum):
    UNKNOWN = 0  # 未知
    HIGHWAY = 1  # Without traffic light / 高速路，无红绿灯
    CITY_EXPRESS = 2  # Without traffic light / 城市快速路，无红绿灯
    CITY_MAIN = 3  # With traffic light / 城市主干道，有红绿灯
    GENERAL_ROAD = 4  # With traffic light / 一般道路，有红绿灯
    COUNTRY_ROAD = 5  # Without traffic light / 乡村道路，无红绿灯


class SpotType(Enum):
    UNKNOWN = 0  # 未知
    TOLL_GATE = 1  # 收费站
    SERVICE_AREA = 2  # 服务区
    TUNNEL = 3  # 隧道
    ACCIDENT = 4  # 事故
    CROSSING = 5  # Intersection / 路口
    ROUNDABOUT = 7  # 环岛
    HIGHWAY_ENTRANCE = 8  # 高速入口
    HIGHWAY_EXIT = 9  # 高速出口
    BRIDGE = 10  # 桥
    EXPRESS_ENTRANCE = 11  # 快速路入口
    EXPRESS_EXIT = 12  # 快速路出口


class Spot:
    def __init__(self):
        self.spot_type: SpotType = SpotType.UNKNOWN  # 场所类型
        self.distance: float = 0.0  # [m] Distance between the spot and subject vehicle / 本车到场所的距离


class TrafficStatus(Enum):
    UNKNOWN = 0  # 未知
    OPEN = 1  # 畅通
    SLOW = 2  # 缓慢
    CROWDED = 3  # 拥挤
    REVERSE = 4  # Only for adjacent lane traffic / 逆向（仅限相邻车道交通状况）


class WindDirection(Enum):
    UNKNOWN = 0  # 未知
    EAST = 1  # 东风
    SOUTH = 2  # 南风
    WEST = 3  # 西风
    NORTH = 4  # 北风
    SOUTH_EAST = 5  # 东南风
    NORTH_EAST = 6  # 东北风
    SOUTH_WEST = 7  # 西南风
    NORTH_WEST = 8  # 西北风


class RoadSurfaceStatus(Enum):
    UNKNOWN = 0  # 未知
    DRY = 1  # 干燥
    WET = 2  # 湿滑
    SNOW_COVERED = 3  # 积雪
    FROZEN = 4  # 结冰


class DisasterStatus(Enum):
    UNKNOWN = 0  # 未知
    NO = 1  # 无灾害
    YES = 2  # 有灾害


class EnvSample:
    def __init__(self):
        self.time: float = 0.0  # [s] Time offset in session / Session内的相对时间
        self.weather_type: WeatherType = WeatherType.UNKNOWN  # 天气类型
        self.temperature: float | None = None  # [°C] 温度
        self.humidity: float | None = None  # [%] Relative humidity / 相对湿度
        self.pressure: float | None = None  # [mb] Atmospheric pressure / 气压
        self.pm25: float | None = None  # [ug/m3] PM 2.5
        self.pm10: float | None = None  # [ug/m3] PM 10
        self.o3_concentration: float | None = None  # [ppb] Ozone concentration / 臭氧浓度
        self.co_concentration: float | None = None  # [ppb] Carbon monoxide concentration / 一氧化碳浓度
        self.so2_concentration: float | None = None  # [ppb] Sulfur dioxide concentration / 二氧化硫浓度
        self.no2_concentration: float | None = None  # [ppb] Nitrogen dioxide concentration / 二氧化氮浓度
        self.visibility: float | None = None  # [KM] 能见度
        self.illumination: float | None = None  # [lux] 光照强度
        self.total_radiation_intensity: float | None = None  # [W/m2] 总辐射强度
        self.ultraviolet_radiation_intensity: float | None = None  # [W/m2] 紫外线辐射强度
        self.rainfall: float | None = None  # [mm/min] Instantaneous rainfall / 瞬时雨量
        self.cumulative_rainfall: float | None = None  # [mm] 累计雨量
        self.wind_speed: float | None = None  # [m/s] 风速
        self.accurate_wind_direction: float | None = None  # [deg] Accurate wind direction (from), 0 is north wind, CCW is positive, ranges -180~180 / 北风为0，CCW为正，范围在-180～180度
        self.wind_direction: WindDirection = WindDirection.UNKNOWN  # Coarse wind direction (from) / 大致风向
        self.typhoon_status: DisasterStatus = DisasterStatus.UNKNOWN  # 台风状况
        self.tornado_status: DisasterStatus = DisasterStatus.UNKNOWN  # 龙卷风状况
        self.lightning_status: DisasterStatus = DisasterStatus.UNKNOWN  # 雷电状况
        self.hail_status: DisasterStatus = DisasterStatus.UNKNOWN  # 冰雹状况
        self.road_name: str | None = None  # 道路名称
        self.speed_limit: int | None = None  # [KPH] 道路限速
        self.road_type: RoadType = RoadType.UNKNOWN  # 道路类型
        self.road_surface_status: RoadSurfaceStatus = RoadSurfaceStatus.UNKNOWN  # 路面状况
        self.road_traffic: TrafficStatus = TrafficStatus.UNKNOWN  # 道路交通状况
        self.current_lane_traffic: TrafficStatus = TrafficStatus.UNKNOWN  # 本车道交通状况
        self.left_lane_traffic: TrafficStatus = TrafficStatus.UNKNOWN  # 左车道交通状况
        self.right_lane_traffic: TrafficStatus = TrafficStatus.UNKNOWN  # 右车道交通状况
        self.spots: List[Spot] = []  # 场所列表

    # Convert to general sample for output / 转通用样本，用于样本输出
    def to_general_sample(self) -> Any:
        output = bi.agency.create_general_sample()
        output.protocol = "env-sample-v3"
        output.time = self.time

        total_length = int(64 + len(self.spots) * 2)
        output.significant = total_length

        values: List[float | str | None] = [None] * total_length
        values[0] = 0
        values[1] = len(self.spots)
        values[2] = self.weather_type.value
        values[3] = self.temperature
        values[4] = self.pm25
        values[5] = self.speed_limit
        values[6] = self.road_type.value
        values[7] = self.illumination
        values[8] = self.wind_speed
        values[9] = self.road_traffic.value
        values[10] = self.road_name
        values[11] = self.pressure
        values[12] = self.humidity
        values[13] = self.visibility
        values[14] = self.wind_direction.value
        values[15] = self.pm10
        values[16] = self.road_surface_status.value
        values[17] = self.current_lane_traffic.value
        values[18] = self.left_lane_traffic.value
        values[19] = self.right_lane_traffic.value
        values[20] = self.typhoon_status.value
        values[21] = self.tornado_status.value
        values[22] = self.lightning_status.value
        values[23] = self.hail_status.value
        values[24] = self.rainfall
        values[25] = self.cumulative_rainfall
        values[26] = self.accurate_wind_direction
        values[27] = self.total_radiation_intensity
        values[28] = self.ultraviolet_radiation_intensity
        values[29] = self.o3_concentration
        values[30] = self.co_concentration
        values[31] = self.so2_concentration
        values[32] = self.no2_concentration

        for i, spot in enumerate(self.spots):
            if isinstance(spot, Spot):
                values[64 + 2 * i] = spot.spot_type.value
                values[65 + 2 * i] = spot.distance
            else:
                values[64 + 2 * i] = SpotType.UNKNOWN.value
                values[65 + 2 * i] = 10000.0
    
        output.values = values
        return output


def _conv_env_sample_v3(gs: Any) -> EnvSample | None:
    values_count = len(gs.values)
    if values_count < 2:
        return None
    road_name_bin_size = int(gs.values[0]) if gs.values[0] is not None else 0
    spot_count = int(gs.values[1]) if gs.values[1] is not None else 0
    if values_count != 64 + 2 * spot_count + road_name_bin_size:
        return None
    output = EnvSample()
    output.time = gs.time
    output.weather_type = WeatherType(int(gs.values[2])) if gs.values[2] is not None else WeatherType.UNKNOWN
    output.temperature = float(gs.values[3]) if gs.values[3] is not None else None
    output.pm25 = float(gs.values[4]) if gs.values[4] is not None else None
    output.speed_limit = int(gs.values[5]) if gs.values[5] is not None else None
    output.road_type = RoadType(int(gs.values[6])) if gs.values[6] is not None else RoadType.UNKNOWN
    output.illumination = float(gs.values[7]) if gs.values[7] is not None else None
    output.wind_speed = float(gs.values[8]) if gs.values[8] is not None else None
    output.road_traffic = TrafficStatus(int(gs.values[9])) if gs.values[9] is not None else TrafficStatus.UNKNOWN
    output.pressure = float(gs.values[11]) if gs.values[11] is not None else None
    output.humidity = float(gs.values[12]) if gs.values[12] is not None else None
    output.visibility = float(gs.values[13]) if gs.values[13] is not None else None
    output.wind_direction = WindDirection(int(gs.values[14])) if gs.values[14] is not None else WindDirection.UNKNOWN
    output.pm10 = float(gs.values[15]) if gs.values[15] is not None else None
    output.road_surface_status = RoadSurfaceStatus(int(gs.values[16])) if gs.values[16] is not None else RoadSurfaceStatus.UNKNOWN
    output.current_lane_traffic = TrafficStatus(int(gs.values[17])) if gs.values[17] is not None else TrafficStatus.UNKNOWN
    output.left_lane_traffic = TrafficStatus(int(gs.values[18])) if gs.values[18] is not None else TrafficStatus.UNKNOWN
    output.right_lane_traffic = TrafficStatus(int(gs.values[19])) if gs.values[19] is not None else TrafficStatus.UNKNOWN
    output.typhoon_status = DisasterStatus(int(gs.values[20])) if gs.values[20] is not None else DisasterStatus.UNKNOWN
    output.tornado_status = DisasterStatus(int(gs.values[21])) if gs.values[21] is not None else DisasterStatus.UNKNOWN
    output.lightning_status = DisasterStatus(int(gs.values[22])) if gs.values[22] is not None else DisasterStatus.UNKNOWN
    output.hail_status = DisasterStatus(int(gs.values[23])) if gs.values[23] is not None else DisasterStatus.UNKNOWN
    output.rainfall = float(gs.values[24]) if gs.values[24] is not None else None
    output.cumulative_rainfall = float(gs.values[25]) if gs.values[25] is not None else None
    output.accurate_wind_direction = float(gs.values[26]) if gs.values[26] is not None else None
    output.total_radiation_intensity = float(gs.values[27]) if gs.values[27] is not None else None
    output.ultraviolet_radiation_intensity = float(gs.values[28]) if gs.values[28] is not None else None
    output.o3_concentration = float(gs.values[29]) if gs.values[29] is not None else None
    output.co_concentration = float(gs.values[30]) if gs.values[30] is not None else None
    output.so2_concentration = float(gs.values[31]) if gs.values[31] is not None else None
    output.no2_concentration = float(gs.values[32]) if gs.values[32] is not None else None
    for i in range(0, spot_count):
        if gs.values[64 + 2 * i] is not None and gs.values[65 + 2 * i] is not None:
            spot = Spot()
            spot.spot_type = SpotType(int(gs.values[64 + 2 * i]))
            spot.distance = float(gs.values[65 + 2 * i])
            output.spots.append(spot)
    if gs.values[10] is not None:
        output.road_name = str(gs.values[10])
    elif road_name_bin_size > 0:
        road_name_bin_ints = []
        for i in range(0, road_name_bin_size):
            b = gs.values[64 + 2 * spot_count + i]
            if b is None:
                road_name_bin_ints = []
                break
            else:
                road_name_bin_ints.append(int(b))
        if len(road_name_bin_ints) > 0:
            output.road_name = bytes(road_name_bin_ints).decode('utf-8')
    return output


def _conv_env_sample_v2(gs: Any) -> EnvSample | None:
    values_count = len(gs.values)
    if values_count < 2:
        return None
    road_name_bin_size = int(gs.values[0]) if gs.values[0] is not None else 0
    spot_count = int(gs.values[1]) if gs.values[1] is not None else 0
    if values_count != 10 + 2 * spot_count + road_name_bin_size:
        return None
    output = EnvSample()
    output.time = gs.time
    output.weather_type = WeatherType(int(gs.values[2])) if gs.values[2] is not None else WeatherType.UNKNOWN
    output.temperature = float(gs.values[3]) if gs.values[3] is not None else None
    output.pm25 = float(gs.values[4]) if gs.values[4] is not None else None
    output.speed_limit = int(gs.values[5]) if gs.values[5] is not None else None
    output.road_type = RoadType(int(gs.values[6])) if gs.values[6] is not None else RoadType.UNKNOWN
    output.illumination = float(gs.values[7]) if gs.values[7] is not None else None
    output.wind_speed = float(gs.values[8]) if gs.values[8] is not None else None
    output.road_traffic = TrafficStatus(int(gs.values[9])) if gs.values[9] is not None else TrafficStatus.UNKNOWN
    for i in range(0, spot_count):
        if gs.values[10 + 2 * i] is not None and gs.values[11 + 2 * i] is not None:
            spot = Spot()
            spot.spot_type = SpotType(int(gs.values[10 + 2 * i]))
            spot.distance = float(gs.values[11 + 2 * i])
            output.spots.append(spot)
    road_name_bin_ints = []
    for i in range(0, road_name_bin_size):
        b = gs.values[10 + 2 * spot_count + i]
        if b is None:
            road_name_bin_ints = []
            break
        else:
            road_name_bin_ints.append(int(b))
    if len(road_name_bin_ints) > 0:
        output.road_name = bytes(road_name_bin_ints).decode('utf-8')
    return output


def _interpolate_env_sample(es1: EnvSample, w1: float, es2: EnvSample, w2: float) -> EnvSample:
    output = EnvSample()
    output.time = bi.time

    if es1.temperature is not None and es2.temperature is not None:
        output.temperature = es1.temperature * w1 + es2.temperature * w2
    if es1.humidity is not None and es2.humidity is not None:
        output.humidity = es1.humidity * w1 + es2.humidity * w2
    if es1.pressure is not None and es2.pressure is not None:
        output.pressure = es1.pressure * w1 + es2.pressure * w2
    if es1.pm25 is not None and es2.pm25 is not None:
        output.pm25 = es1.pm25 * w1 + es2.pm25 * w2
    if es1.pm10 is not None and es2.pm10 is not None:
        output.pm10 = es1.pm10 * w1 + es2.pm10 * w2
    if es1.o3_concentration is not None and es2.o3_concentration is not None:
        output.o3_concentration = es1.o3_concentration * w1 + es2.o3_concentration * w2
    if es1.co_concentration is not None and es2.co_concentration is not None:
        output.co_concentration = es1.co_concentration * w1 + es2.co_concentration * w2
    if es1.so2_concentration is not None and es2.so2_concentration is not None:
        output.so2_concentration = es1.so2_concentration * w1 + es2.so2_concentration * w2
    if es1.no2_concentration is not None and es2.no2_concentration is not None:
        output.no2_concentration = es1.no2_concentration * w1 + es2.no2_concentration * w2
    if es1.visibility is not None and es2.visibility is not None:
        output.visibility = es1.visibility * w1 + es2.visibility * w2
    if es1.illumination is not None and es2.illumination is not None:
        output.illumination = es1.illumination * w1 + es2.illumination * w2
    if es1.total_radiation_intensity is not None and es2.total_radiation_intensity is not None:
        output.total_radiation_intensity = es1.total_radiation_intensity * w1 + es2.total_radiation_intensity * w2
    if es1.ultraviolet_radiation_intensity is not None and es2.ultraviolet_radiation_intensity is not None:
        output.ultraviolet_radiation_intensity = es1.ultraviolet_radiation_intensity * w1 + es2.ultraviolet_radiation_intensity * w2
    if es1.rainfall is not None and es2.rainfall is not None:
        output.rainfall = es1.rainfall * w1 + es2.rainfall * w2
    if es1.cumulative_rainfall is not None and es2.cumulative_rainfall is not None:
        output.cumulative_rainfall = es1.cumulative_rainfall * w1 + es2.cumulative_rainfall * w2
    if es1.wind_speed is not None and es2.wind_speed is not None:
        output.wind_speed = es1.wind_speed * w1 + es2.wind_speed * w2
    if es1.accurate_wind_direction is not None and es2.accurate_wind_direction is not None:
        output.accurate_wind_direction = es1.accurate_wind_direction * w1 + es2.accurate_wind_direction * w2

    output.speed_limit = es1.speed_limit if w1 > w2 else es2.speed_limit
    output.weather_type = es1.weather_type if w1 > w2 else es2.weather_type
    output.road_type = es1.road_type if w1 > w2 else es2.road_type
    output.wind_direction = es1.wind_direction if w1 > w2 else es2.wind_direction
    output.road_surface_status = es1.road_surface_status if w1 > w2 else es2.road_surface_status
    output.road_traffic = es1.road_traffic if w1 > w2 else es2.road_traffic
    output.current_lane_traffic = es1.current_lane_traffic if w1 > w2 else es2.current_lane_traffic
    output.left_lane_traffic = es1.left_lane_traffic if w1 > w2 else es2.left_lane_traffic
    output.right_lane_traffic = es1.right_lane_traffic if w1 > w2 else es2.right_lane_traffic
    output.typhoon_status = es1.typhoon_status if w1 > w2 else es2.typhoon_status
    output.tornado_status = es1.tornado_status if w1 > w2 else es2.tornado_status
    output.lightning_status = es1.lightning_status if w1 > w2 else es2.lightning_status
    output.hail_status = es1.hail_status if w1 > w2 else es2.hail_status
    output.road_name = es1.road_name if w1 > w2 else es2.road_name
    output.spots = es1.spots if w1 > w2 else es2.spots

    return output


# Query EnvSample for input / 获取EnvSample，用于样本输入
def get_env_sample() -> EnvSample | None:
    s1 = None
    s2 = None
    w1 = 0.0
    w2 = 0.0
    protocol_id_v2 = 'env-sample-v2'
    protocol_id_v3 = 'env-sample-v3'
    if protocol_id_v3 in bi.input_samples:
        pair = bi.input_samples[protocol_id_v3]
        s1 = _conv_env_sample_v3(pair.sample1)
        w1 = pair.weight1
        s2 = _conv_env_sample_v3(pair.sample2)
        w2 = pair.weight2
    elif protocol_id_v2 in bi.input_samples:
        pair = bi.input_samples[protocol_id_v2]
        s1 = _conv_env_sample_v2(pair.sample1)
        w1 = pair.weight1
        s2 = _conv_env_sample_v2(pair.sample2)
        w2 = pair.weight2
    if s1 is not None and s2 is not None:
        return _interpolate_env_sample(s1, w1, s2, w2)
    return None
